# -*- coding: utf-8 -*-
"""
中性策略框架 | 邢不行 | 2024分享会
author: 邢不行
微信: xbx6660
"""
import os
import time
import json
import pandas as pd
from glob import glob
from joblib import Parallel, delayed
from datetime import datetime, timedelta
from config import data_center_path, flag_path, utc_offset, data_path


def check_flag(run_time):
    """
    检查flag
    :param run_time:    当前的运行时间
    """
    flag = False
    index_file_path = os.path.join(flag_path, f"{run_time.strftime('%Y-%m-%d_%H')}.flag")  # 构建本地flag文件地址
    while True:
        time.sleep(1)
        # 判断该flag文件是否存在
        if os.path.exists(index_file_path):
            flag = True
            break

        # 当前时间是否超过run_time
        if datetime.now() > run_time + timedelta(minutes=30):  # 如果当前时间超过run_time半小时，表示已经错过当前run_time的下单时间，可能数据中心更新数据失败，没有生成flag文件
            flag = False
            break

    return flag


def read_and_merge_data(file_path, black_list, run_time, cls, is_day_period, get_kline_num, min_kline_size=168):
    """
    读取k线数据，并且合并三方数据
    :param file_path:  k线数据文件
    :param black_list: 黑名单
    :param run_time:   实盘运行时间
    :param cls:        加载的策略脚本
    :param is_day_period: 判断是否是日频率
    :param get_kline_num: 获取多少k线数据
    :param min_kline_size: 最小k线限制
    :return:
    """
    symbol = os.path.basename(file_path).split('.')[0]  # 获取币种名称
    if symbol in black_list:  # 黑名单币种直接跳过
        return symbol, None
    try:
        df = pd.read_csv(file_path, encoding='gbk', parse_dates=['candle_begin_time'])  # 读取k线数据
    except:
        return symbol, None
    df.drop_duplicates(subset=['candle_begin_time'], keep='last', inplace=True)  # 去重保留最新的数据
    df.sort_values('candle_begin_time', inplace=True)  # 通过candle_begin_time排序
    df = df[df['candle_begin_time'] + pd.Timedelta(hours=utc_offset) < run_time]  # 根据run_time过滤一下时间
    if df.shape[0] < min_kline_size:
        return symbol, None

    # 合并数据  跟回测保持一致
    data_dict, factor_dict = {}, {}
    df, factor_dict, data_dict = cls.after_merge_index(df, symbol, factor_dict, data_dict)

    # 转换成日线数据  跟回测保持一致
    if is_day_period:
        df = trans_period_for_day(df, factor_dict=factor_dict)

    df = df[-get_kline_num:]  # 根据config配置，控制内存中币种的数据，可以节约内存，加快计算速度
    df.reset_index(inplace=True, drop=True)  # 重置索引

    return symbol, df


def load_delist():
    """
    加载delist数据，动态处理黑名单
    """
    try:
        de_list_path = os.path.join(data_path, 'delist.json')
        with open(de_list_path, 'r') as file:
            de_list = json.load(file)['list']
        return [_ for _ in de_list if _.endswith('USDT')]  # 获取USDT结尾的币种
    except:
        return []


def load_data(symbol_type, run_time, account_info, njob=1):
    """
    加载数据
    :param symbol_type: 数据类型
    :param run_time:  实盘的运行时间
    :param account_info:  账户配置
    :param njob: 读取数据的核心数
    :return:
    """
    # 加载当前账户下的策略（第一个策略。要求所有策略配置的整体计算 after merge index 的函数都是一致的）
    _cls = __import__(f'strategy.{account_info["strategy_list"][0]["strategy"]}', fromlist=('',))
    # 处理黑名单，兼容回测币种名称
    black_list = [_.replace('-', '') for _ in account_info['black_list']] + load_delist()
    # 获取k线数量。
    get_kline_num = account_info['get_kline_num']
    # 最小k线限制
    min_kline_size = account_info['min_kline_size']
    # 获取策略是日频还是小时频
    is_day_period = account_info['is_day_period']
    # 数据结果集
    result = []

    # 获取当前目录下所有的k线文件路径
    file_list = glob(os.path.join(data_center_path, 'kline', symbol_type, '*.csv'))
    # 根据核心数，决定是否开启多进程
    if njob == 1 or njob == -1:  # -1为全核心运行，为了避免一些进程阻塞，自动变成串行
        for _file in file_list:
            res = read_and_merge_data(_file, black_list, run_time, _cls, is_day_period, get_kline_num, min_kline_size)
            result.append(res)
    else:
        # 多进程获取数据
        result = Parallel(n_jobs=njob)(
            delayed(read_and_merge_data)(_file, black_list, run_time, _cls, is_day_period, get_kline_num, min_kline_size)
            for _file in file_list
        )

    return dict(result)


def trans_period_for_day(df, date_col='candle_begin_time', factor_dict=None):
    """
    将数据周期转换为指定的1D周期
    :param df: 原始数据
    :param date_col: 日期列
    :param factor_dict: 转换规则
    :return:
    """
    df.set_index(date_col, inplace=True)
    # 必备字段
    agg_dict = {
        'symbol': 'first',
        'open': 'first',
        'high': 'max',
        'low': 'min',
        'close': 'last',
        'volume': 'sum',
        'quote_volume': 'sum',
        'symbol_type': 'last',
    }
    if factor_dict:
        agg_dict = dict(agg_dict, **factor_dict)
    df = df.resample('1D').agg(agg_dict)
    df.reset_index(inplace=True)

    return df
